//! Matlab MEX files support
#include "mex.h"
#include "NaiveBayesModel.h"

using namespace Evolutive;
void mexFunction( int nlhs, mxArray *plhs[], int nrhs, const mxArray*prhs[] )
{ 	
	/*
		Entry point. This function is called when the DLL is used in Matlab.

		Parameters:		
			+ nlhs: Number of output arguments.
			+ plhs: Array of output matrix pointers. 
			+ nrhs: Number of input parameters.
			+ prhs: Array of input matrix pointers.
	*/

	double *pData=NULL,*pParameters=NULL,*pAssigment=NULL,*pProb=NULL;
	int NumSamples,DimSamples;
	CNaiveBayesModel Model;
	int NumClusters;
	double *Priors=NULL,*Centres=NULL,*Variances=NULL;
	register int i;


	// Verify the parameters
	if(nrhs<1)
		mexErrMsgTxt("Not enough input parameters");

	// Get the data parameters
	NumSamples=mxGetN(prhs[0]);
	DimSamples=mxGetM(prhs[0]);

	// Point to the data matrix
	pData=mxGetPr(prhs[0]);

	// Configure the model
	if(nrhs>1)
	{
		// Point to the parameters vector
		pParameters=mxGetPr(prhs[1]);

		// Set the parameters
		Model.SetHoldOutFrac(pParameters[0]);
		Model.SetMaxIters(static_cast<int>(pParameters[1]));
		Model.SetMinWeightChange(pParameters[2]);
		Model.SetPruneFrac(pParameters[3]);
		Model.SetPruneFreq(static_cast<int>(pParameters[4]));
		Model.SetInitialNumClusters(static_cast<int>(pParameters[5]));
		Model.SetRemoveClusterCentres(static_cast<bool>(pParameters[6]));
		Model.SetEMDelta(pParameters[7]);
		Model.SetAddDelta(pParameters[8]);
		Model.SetVarDiv(pParameters[9]);
	}

	// Estimate the model
	Model.Estimate(NumSamples,DimSamples,pData);

	// Get the parameters of the model
	NumClusters=Model.GetNumClusters();

	Priors=new double[NumClusters];
	Centres=new double[NumClusters*DimSamples];
	Variances=new double[NumClusters*DimSamples];
				
	Model.GetParameters(Priors,Centres,Variances);

	// Create the output parameters
	// Prior distribution
	plhs[0]=mxCreateDoubleMatrix(1,NumClusters,mxREAL);	
	memcpy(mxGetPr(plhs[0]),Priors,NumClusters*sizeof(double));

	// Centres of the distributions
	if(nlhs>1)
	{
		plhs[1]=mxCreateDoubleMatrix(DimSamples,NumClusters,mxREAL);	
		memcpy(mxGetPr(plhs[1]),Centres,NumClusters*DimSamples*sizeof(double));
	}

	// Covariance matrices
	if(nlhs>2)
	{
		plhs[2]=mxCreateDoubleMatrix(DimSamples,NumClusters,mxREAL);	
		memcpy(mxGetPr(plhs[2]),Variances,NumClusters*DimSamples*sizeof(double));
	}

	if(nlhs>3)
	{
		// Create the assigment matrix
		plhs[3]=mxCreateDoubleMatrix(1,NumSamples,mxREAL);
		pAssigment=mxGetPr(plhs[3]);

		// Create the probability matrix
		if(nlhs>4)
		{
			plhs[4]=mxCreateDoubleMatrix(NumClusters,NumSamples,mxREAL);
			pProb=mxGetPr(plhs[4]);
		}

		// Center assignation
		for(i=0;i<NumSamples;i++)
		{
			if(nlhs>4)
				pAssigment[i]=Model.Classify(&(pData[i*DimSamples]),&(pProb[i*NumClusters]));
			else
				pAssigment[i]=Model.Classify(&(pData[i*DimSamples]));
				
		}
	}

	// Release temporal data
	delete []Priors;
	delete []Centres;
	delete []Variances;
}
